import os
import requests
import m3u8
from concurrent.futures import ThreadPoolExecutor
import subprocess
from pathlib import Path
import tempfile

class M3U8Downloader:
    def __init__(self, max_workers=10):
        self.max_workers = max_workers
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
            'Accept': '*/*',
        })
    
    def download_m3u8_to_video(self, m3u8_url, output_path=None):
        """
        将 m3u8 下载并保存为视频文件
        步骤：
        1. 解析 m3u8
        2. 下载所有 ts 分段
        3. 合并为视频文件
        """
        # 1. 解析 m3u8
        print("步骤1: 解析 m3u8 文件...")
        m3u8_obj = m3u8.load(m3u8_url)
        
        if not m3u8_obj.segments:
            print("错误: 未找到视频分段")
            return False
        
        print(f"找到 {len(m3u8_obj.segments)} 个视频分段")
        
        # 设置输出路径
        if output_path is None:
            download_dir = str(Path.home() / "Downloads")
            output_path = os.path.join(download_dir, "downloaded_video.mp4")
        
        # 2. 下载所有 ts 分段
        print("步骤2: 下载视频分段...")
        ts_files = self.download_all_segments(m3u8_obj, m3u8_url)
        
        if not ts_files:
            print("错误: 没有成功下载任何分段")
            return False
        
        # 3. 合并为视频文件
        print("步骤3: 合并为视频文件...")
        success = self.merge_to_video(ts_files, output_path)
        
        if success:
            print(f"✅ 视频保存成功: {output_path}")
        else:
            print("❌ 视频合并失败")
        
        return success
    
    def download_all_segments(self, m3u8_obj, m3u8_url):
        """下载所有 ts 分段"""
        # 获取基础 URL（用于相对路径）
        base_url = self.get_base_url(m3u8_url)
        
        # 创建临时目录
        temp_dir = tempfile.mkdtemp()
        ts_files = []
        
        def download_segment(segment_info):
            """下载单个分段"""
            index, segment = segment_info
            try:
                # 构建完整的 URL
                if segment.uri.startswith(('http://', 'https://')):
                    url = segment.uri
                else:
                    url = base_url + segment.uri
                
                response = self.session.get(url, timeout=30)
                response.raise_for_status()
                
                # 保存 ts 文件
                ts_path = os.path.join(temp_dir, f"segment_{index:05d}.ts")
                with open(ts_path, 'wb') as f:
                    f.write(response.content)
                
                print(f"✅ 已下载分段 {index + 1}/{len(m3u8_obj.segments)}")
                return ts_path
                
            except Exception as e:
                print(f"❌ 下载分段 {index} 失败: {e}")
                return None
        
        # 使用多线程下载
        segment_infos = [(i, seg) for i, seg in enumerate(m3u8_obj.segments)]
        
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            results = executor.map(download_segment, segment_infos)
        
        # 收集成功的文件
        ts_files = [result for result in results if result is not None]
        
        return ts_files
    
    def get_base_url(self, m3u8_url):
        """获取基础 URL"""
        if m3u8_url.startswith(('http://', 'https://')):
            return '/'.join(m3u8_url.split('/')[:-1]) + '/'
        return ''
    
    def merge_to_video(self, ts_files, output_path):
        """将 ts 文件合并为视频文件"""
        if not ts_files:
            return False
        
        try:
            # 方法1: 使用 ffmpeg 合并（推荐）
            return self.merge_with_ffmpeg(ts_files, output_path)
        except:
            # 方法2: 直接文件合并
            return self.merge_direct(ts_files, output_path)
    
    def merge_with_ffmpeg(self, ts_files, output_path):
        """使用 ffmpeg 合并为 MP4"""
        try:
            # 创建文件列表
            list_file = output_path + '_filelist.txt'
            with open(list_file, 'w', encoding='utf-8') as f:
                for ts_file in sorted(ts_files):  # 确保顺序正确
                    f.write(f"file '{os.path.abspath(ts_file)}'\n")
            
            # ffmpeg 命令
            cmd = [
                'ffmpeg',
                '-f', 'concat',
                '-safe', '0',
                '-i', list_file,
                '-c', 'copy',  # 直接复制，不重新编码
                output_path,
                '-y'
            ]
            
            print("使用 ffmpeg 合并视频...")
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            # 清理临时文件
            if os.path.exists(list_file):
                os.remove(list_file)
            
            # 清理 ts 文件
            for ts_file in ts_files:
                if os.path.exists(ts_file):
                    os.remove(ts_file)
            if ts_files and os.path.exists(os.path.dirname(ts_files[0])):
                os.rmdir(os.path.dirname(ts_files[0]))
            
            return result.returncode == 0
            
        except Exception as e:
            print(f"ffmpeg 合并失败: {e}")
            return False
    
    def merge_direct(self, ts_files, output_path):
        """直接合并 ts 文件"""
        try:
            # 保存为 .ts 文件
            ts_output = output_path.replace('.mp4', '.ts')
            
            with open(ts_output, 'wb') as outfile:
                for ts_file in sorted(ts_files):
                    with open(ts_file, 'rb') as infile:
                        outfile.write(infile.read())
            
            print(f"TS 文件已保存: {ts_output}")
            print("提示: 可以使用 ffmpeg 转换为 MP4: ffmpeg -i input.ts -c copy output.mp4")
            
            # 清理临时文件
            for ts_file in ts_files:
                if os.path.exists(ts_file):
                    os.remove(ts_file)
            if ts_files and os.path.exists(os.path.dirname(ts_files[0])):
                os.rmdir(os.path.dirname(ts_files[0]))
            
            return True
            
        except Exception as e:
            print(f"直接合并失败: {e}")
            return False

# 使用示例
if __name__ == "__main__":
    downloader = M3U8Downloader()
    
    # m3u8 链接
    m3u8_url = ""  # 替换为实际的 m3u8 URL
    
    # 下载并保存为视频
    success = downloader.download_m3u8_to_video(m3u8_url)
    
    if success:
        print("🎉 下载完成！")
    else:
        print("💥 下载失败！")